import matplotlib.pyplot as plt
import numpy as np
import matplotlib
import pandas as pd
import matplotlib.pyplot as pyplot
import sklearn
from sklearn import tree, datasets, preprocessing
from matplotlib import style

# 获取数据
data = pd.read_csv("../winequality-white.csv", sep=";")
# print(data.head())
# 获取数据中这五个属性
data = data[["total sulfur dioxide", "density", "pH", "sulphates", "alcohol", "quality"]]

predict = "quality"

x = np.array(data.drop([predict], axis=1))
# 以G3作为因变量
y = np.array(data[predict])

x_train, x_test, y_train, y_test = sklearn.model_selection.train_test_split(x, y, test_size=0.1)

# clf = tree.DecisionTreeRegressor(random_state=666, splitter='random', max_depth=2)
clf = tree.DecisionTreeClassifier(random_state=666, splitter='random', max_depth=2)
clf = clf.fit(x_train, y_train)
result = clf.score(x_test, y_test)
print(result)
predicted = clf.predict(x_test)
for x in range(len(predicted)):
    print(predicted[x], x_test[x], y_test[x])
p = "pH"  # 这里的p可以是任何属性，通过改变这里的p即可获得该属性与Final Grade之间的关系
style.use("ggplot")
pyplot.scatter(data[p], data["quality"])
pyplot.xlabel(p)
pyplot.ylabel("Final Class")
pyplot.show()
